/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "parcsr_blockjacobisor_preconditioner.h"

#ifdef USE_OPENMP
#include <omp.h>
#endif

namespace YHAMG
{

ParCSR_BlockJacobiSOR_Preconditioner::ParCSR_BlockJacobiSOR_Preconditioner()
	: num_threads(0),
	Type(2),
	Factor(1.0)
{
}

ParCSR_BlockJacobiSOR_Preconditioner::ParCSR_BlockJacobiSOR_Preconditioner(int type, int factor)
	: num_threads(0),
	Type(type),
	Factor(factor)
{
}

void ParCSR_BlockJacobiSOR_Preconditioner::Free()
{
	num_threads = 0;
	D.Free();
	D_Recip.Free();
	L.Free();
	U.Free();
}

void ParCSR_BlockJacobiSOR_Preconditioner::Setup(const ParCSR_Matrix& A, int REUSE)
{
	int n = A.LocLoc.size[0];
	int* Ap = A.LocLoc.rowptr;
	int* Ai = A.LocLoc.colind;
	double* Av = A.LocLoc.values;
	
	if (REUSE == 0)
	{
		comm = A.comm;

		Free();

		D.Resize(n);
		D_Recip.Resize(n);

		double* Dv = D.values;
		double* Dv_Recip = D_Recip.values;

		num_threads = 1;
#ifdef USE_OPENMP
		num_threads = omp_get_max_threads();
#endif

		int* Lp = new int[n + 1];
		int* Up = new int[n + 1];

		Lp[0] = 0;
		Up[0] = 0;
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
		for (int t = 0; t < num_threads; ++t)
		{
			int begin = t * n / num_threads;
			int end = (t + 1) * n / num_threads;

			for (int i = begin; i < end; ++i)
			{
				int cnt_l = 0;
				int cnt_u = 0;

				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
				{
					int jcol = Ai[j];

					if (jcol < begin || jcol >= end) continue;

					if (jcol < i)
						++cnt_l;
					else if (jcol > i)
						++cnt_u;
				}

				Lp[i + 1] = cnt_l;
				Up[i + 1] = cnt_u;
			}
		}

		for (int i = 0; i < n;++i)
			Lp[i + 1] += Lp[i];
		for (int i = 0; i < n;++i)
			Up[i + 1] += Up[i];

		int* Li = new int[Lp[n]];
		double* Lv = new double[Lp[n]];
		int* Ui = new int[Up[n]];
		double* Uv = new double[Up[n]];

#ifdef USE_OPENMP
#pragma omp parallel for
#endif
		for (int t = 0; t < num_threads; ++t)
		{
			int begin = t * n / num_threads;
			int end = (t + 1) * n / num_threads;

			for (int i = begin; i < end; ++i)
			{
				for (int j = Ap[i], k = Lp[i], r = Up[i]; j < Ap[i + 1]; ++j)
				{
					int jcol = Ai[j];
					double jval = Av[j];

					if (jcol < begin || jcol >= end) continue;

					if (jcol < i)
					{
						Li[k] = jcol;
						Lv[k++] = jval;
					}
					else if (jcol > i)
					{
						Ui[r] = jcol;
						Uv[r++] = jval;
					}
					else
					{
						Dv[i] = jval;
						Dv_Recip[i] = 1.0 / jval;
					}
				}
			}
		}

		L.size[0] = n;
		L.size[1] = n;
		L.rowptr = Lp;
		L.colind = Li;
		L.values = Lv;
		U.size[0] = n;
		U.size[1] = n;
		U.rowptr = Up;
		U.colind = Ui;
		U.values = Uv;
	}
	else
	{
		int* Lp = L.rowptr;
		int* Li = L.colind;
		double* Lv = L.values;
		int* Up = U.rowptr;
		int* Ui = U.colind;
		double* Uv = U.values;
		double* Dv = D.values;
		double* Dv_Recip = D_Recip.values;
		
		int* w = new int[n];

		for (int i = 0; i < n; ++i)
			w[i] = -1;

#ifdef USE_OPENMP
#pragma omp parallel for
#endif
		for (int t = 0; t < num_threads; ++t)
		{
			int begin = t * n / num_threads;
			int end = (t + 1) * n / num_threads;

			for (int i = begin; i < end; ++i)
			{
				int j0 = Ap[i];

				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					if (Ai[j] >= begin && Ai[j] < end) w[Ai[j]] = j;

				if (w[i] >= j0)
				{
					Dv[i] = Av[w[i]];
					Dv_Recip[i] = 1.0 / Av[w[i]];
				}

				for (int j = Lp[i]; j < Lp[i + 1]; ++j)
					if (w[Li[j]] >= j0) Lv[j] = Av[w[Li[j]]];
				for (int j = Up[i]; j < Up[i + 1]; ++j)
					if (w[Ui[j]] >= j0) Uv[j] = Av[w[Ui[j]]];
			}
		}
		
		delete[] w;	
	}
}

int ParCSR_BlockJacobiSOR_Preconditioner::InSize() const
{
	return D.size;
}

int ParCSR_BlockJacobiSOR_Preconditioner::OutSize() const
{
	return D.size;
}

void ParCSR_BlockJacobiSOR_Preconditioner::Apply(const Par_Vector& b, const Par_Vector& x) const
{
	int type = Type;
	double omega = Factor;
	
	int n = D.size;
	double* Dv = D.values;
	double* Dv_Recip = D_Recip.values;
	int* Lp = L.rowptr;
	int* Li = L.colind;
	double* Lv = L.values;
	int* Up = U.rowptr;
	int* Ui = U.colind;
	double* Uv = U.values;
	double* bv = b.Loc.values;
	double* xv = x.Loc.values;

	if (n == 0) return;

	for (int i = 0; i < n; ++i)
		xv[i] = bv[i];

#ifdef USE_OPENMP
#pragma omp parallel for
#endif
	for (int t = 0; t < num_threads; ++t)
	{
		int begin = t * n / num_threads;
		int end = (t + 1) * n / num_threads;

		if (type == 0 || type == 2)
		{
			for (int i = begin; i < end; ++i)
			{
				double temp = 0.0;
				for (int j = Lp[i]; j < Lp[i + 1]; ++j)
					temp -= xv[Li[j]] * Lv[j];
				xv[i] = (xv[i] + omega * temp) * Dv_Recip[i];
			}
		}

		if (type == 2)
		{
			for (int i = begin; i < end; ++i)
				xv[i] *= Dv[i];
		}

		if (type == 1 || type == 2)
		{
			for (int i = end - 1; i >= begin; --i)
			{
				double temp = 0.0;
				for (int j = Up[i + 1] - 1; j >= Up[i]; --j)
					temp -= xv[Ui[j]] * Uv[j];
				xv[i] = (xv[i] + omega * temp) * Dv_Recip[i];
			}
		}
	}
}

}